package lzc.encrypt.encryption;

import lzc.encrypt.Assit;

import java.math.BigInteger;
import java.util.Random;
import java.util.concurrent.ThreadLocalRandom;

/**
 * @author 蓝色的三明治
 * @date 2024/10/12
 */
public class BigIntMatrix {
    public static int[] limit_matrix = new int[]{10, 50};
    public static int[] limit_calculate = new int[]{50, 100};

    private BigInteger[][] matrix;
    private int size;  // 方阵的大小

    /**
     * 生成一个逆矩阵是整数矩阵的随机整数矩阵（建议控制 size 大小，否则程序将陷入很长很长的时间用于计算）
     * 通过计算多个满足条件的 0-1 矩阵，互乘获得满足条件的矩阵
     *
     * @param size        矩阵大小
     * @param lower_limit 整数下限
     * @param upper_limit 整数上限
     */
    public BigIntMatrix(int size, BigInteger lower_limit, BigInteger upper_limit) {
        this.size = size;
        this.matrix = new BigInteger[size][size];

        if (lower_limit.equals(BigInteger.ZERO) && upper_limit.equals(BigInteger.ONE)) {
            // 构造逆矩阵是整数矩阵的随机 0-1 矩阵，直接随机生成 0-1 矩阵直至满足条件即可
            do {
                generateRandomMatrix(lower_limit, upper_limit);
            } while (!isInverseInt());
        } else {
            // 生成的随机整数可逆矩阵，需要满足逆矩阵同样也是整数矩阵的要求
            // 直接采用循环生成随机整数矩阵，直至该矩阵的逆矩阵也是整数矩阵时退出循环是最简单的方法，但是随机整数矩阵要满足该条件几乎不可能
            // 所以，该方法采用，随机生成若干个 0-1 矩阵，这些矩阵的逆矩阵是整数矩阵，若干次互乘得到一个结果矩阵，该结果矩阵的逆矩阵必定是整数矩阵
            // 显然，随机生成的 0-1 矩阵，满足逆矩阵是整数矩阵的条件要容易得多，但 size 过大时同样难以实现，建议控制 size 大小
            do {
                // matrix 初始化为单位矩阵
                for (int i = 0; i < size; i++) {
                    for (int j = 0; j < size; j++) {
                        this.matrix[i][j] = i == j ? BigInteger.ONE : BigInteger.ZERO;
                    }
                }

                int num_matrix = ThreadLocalRandom.current().nextInt(limit_matrix[0], limit_matrix[1]);
                int num_calculate = ThreadLocalRandom.current().nextInt(limit_calculate[0], limit_calculate[1]);

                // 计算 num_matrix 个 0-1 矩阵，用于计算最后的整数矩阵
                BigIntMatrix[] middle_matrix = new BigIntMatrix[num_matrix];
                for (int i = 0; i < num_matrix; i++) {
                    middle_matrix[i] = new BigIntMatrix(size, BigInteger.ZERO, BigInteger.ONE);
                }

                // 进行 num_calculate 次互乘运算，得到一个逆矩阵是整数矩阵的随机整数矩阵
                for (int num = 0; num < num_calculate; num++) {
                    int id_middle_matrix = ThreadLocalRandom.current().nextInt(0, num_matrix);
                    this.matrix = this.times(middle_matrix[id_middle_matrix]).getMatrix();
                }
            } while (!isInRange(lower_limit, upper_limit));
        }
    }

    /**
     * 构造方法
     *
     * @param size 矩阵大小
     */
    public BigIntMatrix(int size) {
        this.size = size;
        this.matrix = new BigInteger[size][size];
    }


    /**
     * 随机生成整数矩阵
     *
     * @param lower_limit 整数下限
     * @param upper_limit 整数上限
     */
    private void generateRandomMatrix(BigInteger lower_limit, BigInteger upper_limit) {
        Random random = new Random();
        for (int i = 0; i < size; i++) {
            for (int j = 0; j < size; j++) {
                // 生成 1 到 upper_limit 范围内的随机整数
                this.matrix[i][j] = Assit.getRandom(lower_limit, upper_limit);
            }
        }
    }

    // 计算矩阵的行列式
    private BigInteger determinant() {
        return computeDeterminant(matrix);
    }

    // 递归计算行列式
    private BigInteger computeDeterminant(BigInteger[][] mat) {
        int n = mat.length;
        if (n == 1) {
            return mat[0][0];
        }
        if (n == 2) {
            return mat[0][0].multiply(mat[1][1]).subtract(mat[0][1].multiply(mat[1][0]));
        }

        BigInteger det = BigInteger.ZERO;
        for (int col = 0; col < n; col++) {
            det = det.add(mat[0][col].multiply(cofactor(mat, 0, col)));
        }
        return det;
    }

    // 计算余子式
    private BigInteger cofactor(BigInteger[][] mat, int row, int col) {
        BigInteger[][] minor = getMinor(mat, row, col);
        BigInteger sign = (row + col) % 2 == 0 ? BigInteger.ONE : BigInteger.valueOf(-1);
        return sign.multiply(computeDeterminant(minor));
    }

    // 获取余子式
    private BigInteger[][] getMinor(BigInteger[][] mat, int row, int col) {
        int n = mat.length;
        BigInteger[][] minor = new BigInteger[n - 1][n - 1];
        int r = 0;
        for (int i = 0; i < n; i++) {
            if (i == row) continue;
            int c = 0;
            for (int j = 0; j < n; j++) {
                if (j == col) continue;
                minor[r][c] = mat[i][j];
                c++;
            }
            r++;
        }
        return minor;
    }

    /**
     * 判断矩阵是否可逆（行列式不为0）
     *
     * @return boolean
     */
    private boolean isInvertible() {
        BigInteger det = determinant();
        return !det.equals(BigInteger.ZERO);
    }

    /**
     * 判断矩阵的逆矩阵是否是整数矩阵（行列式为1或-1）
     *
     * @return boolean
     */
    private boolean isInverseInt() {
        BigInteger det = determinant();
        return det.equals(BigInteger.ONE) || det.equals(BigInteger.valueOf(-1));
    }

    /**
     * 判断矩阵的值是否在指定范围内
     *
     * @param lower_limit 整数下限
     * @param upper_limit 整数上限
     * @return boolean true表示在指定范围内，false表示不在指定范围内
     */
    private boolean isInRange(BigInteger lower_limit, BigInteger upper_limit) {
        for (int i = 0; i < size; i++) {
            for (int j = 0; j < size; j++) {
                if (this.matrix[i][j].compareTo(lower_limit) < 0 || this.matrix[i][j].compareTo(upper_limit) > 0) {
                    return false;
                }
            }
        }
        return true;
    }

    // 计算矩阵的逆矩阵
    public BigIntMatrix inverse() throws IllegalArgumentException {
        BigInteger det = determinant();

        BigInteger[][] adjugate = adjugate(matrix);
        BigInteger[][] inverse = new BigInteger[size][size];

        for (int i = 0; i < size; i++) {
            for (int j = 0; j < size; j++) {
                inverse[i][j] = adjugate[i][j].divide(det);  // 这里除以1或-1保持整数
            }
        }

        return new BigIntMatrix(inverse);
    }

    // 计算伴随矩阵
    private BigInteger[][] adjugate(BigInteger[][] mat) {
        BigInteger[][] adj = new BigInteger[size][size];
        for (int i = 0; i < size; i++) {
            for (int j = 0; j < size; j++) {
                adj[j][i] = cofactor(mat, i, j);  // 转置余子式矩阵
            }
        }
        return adj;
    }

    // 显示矩阵
    public void display() {
        for (int i = 0; i < size; i++) {
            for (int j = 0; j < size; j++) {
                System.out.print(matrix[i][j] + " ");
            }
            System.out.println();
        }
    }

    // 私有构造方法，用于生成逆矩阵
    private BigIntMatrix(BigInteger[][] matrix) {
        this.size = matrix.length;
        this.matrix = matrix;
    }

    /**
     * 矩阵与向量相乘
     *
     * @param vector 整数向量
     * @return {@link BigInteger[]} 相乘后的结果向量
     * @throws IllegalArgumentException
     */
    public BigInteger[] times(BigInteger[] vector) throws IllegalArgumentException {
        if (vector.length != size) {
            throw new IllegalArgumentException("向量长度必须等于矩阵的大小！");
        }
        BigInteger[] result = new BigInteger[size];
        for (int i = 0; i < size; i++) {
            result[i] = BigInteger.ZERO;
            for (int j = 0; j < size; j++) {
                result[i] = result[i].add(matrix[i][j].multiply(vector[j]));// 矩阵乘法运算
            }
        }
        return result;
    }

    /**
     * 向量与矩阵相乘
     *
     * @param vector 整数向量
     * @return {@link BigInteger[]} 相乘后的结果向量
     * @throws IllegalArgumentException
     */
    public BigInteger[] times4(BigInteger[] vector) throws IllegalArgumentException {
        if (vector.length != size) {
            throw new IllegalArgumentException("向量长度必须等于矩阵的大小！");
        }
        BigInteger[] result = new BigInteger[size];
        for (int j = 0; j < size; j++) {
            result[j] = BigInteger.ZERO;
            for (int i = 0; i < size; i++) {
                result[j] = result[j].add(vector[i].multiply(matrix[i][j]));// 矩阵乘法运算
            }
        }
        return result;
    }

    // 设定矩阵中的元素
    public void setElement(int row, int col, BigInteger value) {
        this.matrix[row][col] = value;
    }

    // 获取矩阵中的元素
    public BigInteger getElement(int row, int col) {
        return this.matrix[row][col];
    }

    // 矩阵相乘方法
    public BigIntMatrix times(BigIntMatrix other) {
        if (this.size != other.size) {
            throw new IllegalArgumentException("矩阵大小必须相同才能相乘");
        }

        BigIntMatrix result = new BigIntMatrix(this.size);

        // 矩阵相乘规则：C = A * B
        for (int i = 0; i < size; i++) {
            for (int j = 0; j < size; j++) {
                BigInteger sum = BigInteger.ZERO;
                for (int k = 0; k < size; k++) {
                    sum = sum.add(this.matrix[i][k].multiply(other.matrix[k][j]));
                }
                result.setElement(i, j, sum);
            }
        }
        return result;
    }

    public BigInteger[][] getMatrix() {
        return matrix;
    }
}
